#include "stdafx.h"
#include "TestFunctionHandler.h"

CTestFunctionHandler::CTestFunctionHandler(const EnvDTE::CodeElementPtr & spTestGroupElement, CAddTracker * pAddTracker) :
	m_spTestGroup(spTestGroupElement), m_pAddTracker(pAddTracker)
{

}

CTestFunctionHandler::~CTestFunctionHandler(void)
{
}

BOOL CTestFunctionHandler::Add(LPCWSTR lpTemplate)
{
	BOOL bRet = FALSE;
	if(DefineTestFunctionName(lpTemplate))
	{
		bRet =	DefineUtFuncTable();

		//for a new created test group there could not be information 
		// in code model so let's give time devend to parse sources
		for (int n = 0; (n < 100) && (bRet == FALSE); ++n)
		{
			Sleep(100);
			bRet =	DefineUtFuncTable();

		}

		if ((bRet) &&
			AddEntryInUtFuncTable() &&
			AddNewTestFunction())
		{
			bRet = TRUE;
		}
	}
	return bRet;
}

void CTestFunctionHandler::Add(CElementTree<CodeTreeItem> & rNode)
{
	BOOL bRet =	DefineUtFuncTable();

	//for a new created test group there could not be information 
	// in code model so let's give time devenv to parse sources
	for (int n = 0; (n < 100) && (bRet == FALSE); ++n)
	{
		Sleep(100);
		bRet =	DefineUtFuncTable();
	}

	CAtlStringW wstrTableEntry;

	if (rNode.GetValue().Checked)
	{
		DefineTestFunctionName(L"Test");
		AddNewTestFunction();

		wstrTableEntry = L"\tUT_FUNC_ENTRY(";
		wstrTableEntry.Append(m_bstrTestFunctionName);
		wstrTableEntry.Append(L")\r\n");
	}

	POSITION pos = rNode.GetHeadPosition();
	while(pos)
	{
		const CodeTreeItem & rValue = rNode.GetNext(pos).GetValue();
		if ((rValue.Checked) && (rValue.eKind == EnvDTE::vsCMElementFunction))
		{
			CAtlStringW wstrFuncName(rValue.Name.GetString());
			wstrFuncName.Append(L"Test");

			if (rValue.ElementData.eFunctionKind & EnvDTE::vsCMFunctionConstructor)
				wstrFuncName = L"ConstructorTest";
			else if (rValue.ElementData.eFunctionKind & EnvDTE::vsCMFunctionDestructor)
				wstrFuncName = L"DestructorTest";
			else if (rValue.ElementData.eFunctionKind & EnvDTE::vsCMFunctionOperator)
				wstrFuncName = L"OperatorTest";

			DefineTestFunctionName(wstrFuncName);
			AddNewTestFunction();

			wstrTableEntry.Append(L"\tUT_FUNC_ENTRY(");
			wstrTableEntry.Append(m_bstrTestFunctionName);
			wstrTableEntry.Append(L")\r\n");
		}
	}

	if (wstrTableEntry.GetLength())
	{
		AddEntryInUtFuncTable(wstrTableEntry);
	}
}


BOOL CTestFunctionHandler::DefineTestFunctionName(LPCWSTR lpTemplate)
{
	BOOL bRet = FALSE;
	long nCount = 0;
	EnvDTE::CodeElementsPtr spCodeElements;
	if (SUCCEEDED(m_spTestGroup->get_Functions(&spCodeElements)) &&
		SUCCEEDED(spCodeElements->get_Count(&nCount)))
	{
		CComBSTR bstr;
		EnvDTE::CodeElementPtr spCodeElement;
		CAtlMap<CAtlStringW, DWORD> mapFunctionNames; //actually is used as set
		for (CComVariant vt(1); vt.lVal <= nCount; ++vt.lVal)
		{
			if (SUCCEEDED(spCodeElements->Item(vt, &spCodeElement)) &&
				SUCCEEDED(spCodeElement->get_Name(&bstr)))
			{
				mapFunctionNames.SetAt(bstr, 1);
			}
		}

		CAtlStringW wstrTemplate = lpTemplate;
		for (int n = 0; n < wstrTemplate.GetLength(); ++n)
		{
			if (false == iswalnum(wstrTemplate[n]))
				wstrTemplate.SetAt(n, L'_');
		}


		CAtlStringW wstr = wstrTemplate;
		int nUniqueId = 0;
		do
		{
			if (NULL == mapFunctionNames.Lookup(wstr))
			{
				m_bstrTestFunctionName = wstr;
				bRet = TRUE;
				break;
			}

			++nUniqueId;
			wstr = wstrTemplate;
			wstr.AppendFormat(L"%i", nUniqueId);
		}
		while (nUniqueId < 10000);
	}

	return bRet;
}


BOOL CTestFunctionHandler::DefineUtFuncTable()
{
	BOOL bRet = FALSE;
	EnvDTE::CodeElementsPtr spCodeElements;
	CComBSTR bstrName;
	EnvDTE::CodeElementPtr spElement;
	long nCount = 0;

	if (SUCCEEDED(m_spTestGroup->get_FullName(&bstrName)) &&
		SUCCEEDED(m_spTestGroup->get_Functions(&spCodeElements)) &&
		SUCCEEDED(spCodeElements->get_Count(&nCount)))
	{
		bstrName.Append(L"::GetTestsFunc");
		for (CComVariant vt(1); vt.lVal <= nCount; ++vt.lVal)
		{
			CComBSTR bstr;
			if (SUCCEEDED(spCodeElements->Item(vt, &spElement)) &&
				SUCCEEDED(spElement->get_FullName(&bstr)) &&
				(bstr == bstrName))
			{
				EnvDTE::ProjectItemPtr spProjItem;
				short nFileCount = 0;
				if (SUCCEEDED(spElement->get_StartPoint(&m_spUtFuncTableStart)) &&
					SUCCEEDED(spElement->get_EndPoint(&m_spUtFuncTableEnd)) &&
					SUCCEEDED(spElement->get_ProjectItem(&spProjItem)) &&
					SUCCEEDED(spProjItem->get_FileCount(&nFileCount)) &&
					(nFileCount) &&
					SUCCEEDED(spProjItem->get_FileNames(0, &m_bstrSourceFilePath)))
				{
					bRet = TRUE;
				}
				break;
			}
		}
	}

	return bRet; //not found in case UtTestBase.h is no visible from the test project
}

BOOL CTestFunctionHandler::AddNewTestFunction()
{
	BOOL bRet = FALSE;
	EnvDTE::CodeFunctionPtr spCodeFunction;
	VCCodeModelLibrary::VCCodeFunctionPtr spNewTestFunc;
	
	if (SUCCEEDED(m_spTestGroup->AddFunction(m_bstrTestFunctionName, EnvDTE::vsCMFunctionFunction,
		CComVariant(EnvDTE::vsCMTypeRefVoid), CComVariant(-1), 
		EnvDTE::vsCMAccessPublic, CComVariant(m_bstrSourceFilePath), &spCodeFunction)) &&
		(spNewTestFunc = spCodeFunction) &&
		SUCCEEDED(spNewTestFunc->put_BodyText(CComBSTR(L"UT_ASSERT2(FALSE, L\"not implemented\");\r\n"))))
	{
		bRet = TRUE;
		if (m_pAddTracker)
			m_pAddTracker->AddTestFunc(spNewTestFunc);
	}

	return bRet;
}

BOOL CTestFunctionHandler::AddEntryInUtFuncTable()
{
	CAtlStringW wstrTableEntry = L"\tUT_FUNC_ENTRY(";
	wstrTableEntry.Append(m_bstrTestFunctionName);
	wstrTableEntry.Append(L")\r\n");
	
	return AddEntryInUtFuncTable(wstrTableEntry);
}

BOOL CTestFunctionHandler::AddEntryInUtFuncTable(LPCWSTR lpTableContent)
{
	BOOL bRet = FALSE;
	CComBSTR bstr;
	CComBSTR bstrGroupName;
	EnvDTE::EditPointPtr spPt;

	CAtlStringW wstrTableEntry = L"\tUT_FUNC_ENTRY(";
	wstrTableEntry.Append(m_bstrTestFunctionName);
	wstrTableEntry.Append(L")\r\n");
	
	long nStart = 0, nEnd = 0;

	if (SUCCEEDED(m_spUtFuncTableStart->CreateEditPoint(&spPt)) &&
		SUCCEEDED(m_spUtFuncTableStart->get_Line(&nStart)) &&
		SUCCEEDED(m_spUtFuncTableEnd->get_Line(&nEnd)))
	{
		for(long n = nStart; n <= nEnd; ++n)
		{
			if (SUCCEEDED(spPt->LineDown(1)) &&
				SUCCEEDED(spPt->GetLines(n, n + 1, &bstr)))
			{
				CAtlStringW wstr = bstr;
				if (-1 != wstr.Find(L"END_UT_FUNC_TABLE")) //don't care if it is comment
				{
					if (SUCCEEDED(spPt->LineUp(1)) &&
						SUCCEEDED(spPt->Insert(CComBSTR(lpTableContent))))
					{
						bRet = TRUE;
					}
					break;
				}
			}
		}
	}
	return bRet;
}